package com.tuniu.agents.common.handler;

import com.tuniu.agents.advisor.ChatEnvironmentMemory;
import com.tuniu.agents.common.command.CommandManagerUtil;
import com.tuniu.agents.common.command.action.DefaultCallAgentAction;
import com.tuniu.agents.common.command.action.DefaultOutputTextAction;
import com.tuniu.agents.common.command.match.StringMatch;
import com.tuniu.agents.common.entity.itinerary.BaseItinerary;
import com.tuniu.agents.common.entity.itinerary.Resource;
import com.tuniu.agents.common.entity.itinerary.enums.ResourceStatus;
import com.tuniu.agents.common.entity.itinerary.resource.TrainResource;
import com.tuniu.agents.common.i18n.I18nConstants;
import com.tuniu.agents.common.message.BaseAgentNewResp;
import com.tuniu.agents.common.message.CardTypeEnum;
import com.tuniu.agents.common.service.itinerary.ItineraryService;
import com.tuniu.agents.common.util.NumberingUtil;
import com.tuniu.agents.impl.train.TrainAgentConfig;
import com.tuniu.agents.memory.MemoryManager;
import com.tuniu.agents.message.MessageRender;
import com.tuniu.agents.tools.train.model.TrainCardData;
import com.tuniu.agents.tools.train.tool.SearchAllTrainTool;
import com.tuniu.agents.tools.train.tool.SearchTrainDetailTool;
import com.tuniu.agents.util.I18nUtil;
import com.tuniu.agents.util.JsonUtil;
import lombok.extern.slf4j.Slf4j;
import org.apache.groovy.util.Maps;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.math.BigDecimal;
import java.util.*;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

import static com.tuniu.agents.tools.train.tool.SearchAllTrainTool.buildMarkdownAssistantMsg;
import static com.tuniu.agents.tools.train.utils.TrainCommonUtil.extractDate;


/**
 * Train ticket streaming response handler implementation
 */
@Slf4j
public class TrainStreamResponseHandler implements IStreamResponseHandler {

    private final String questionId;

    private final String toolName;

    private final String agentId;
    private final ChatEnvironmentMemory chatEnvironmentMemory;

    private final MemoryManager memoryManager;

    // Matching patterns
    private final Pattern bracketPattern;

    private final Pattern angleBracketPattern = Pattern.compile("<([^>]+)>");

    public static final String MESSAGE = "message";

    public static final String CONVERSATION_ID = "conversationId";

    // Pattern for matching train numbers starting with letters followed by numbers (e.g., G2578)
    private final Pattern trainNumberPattern = Pattern.compile("^([A-Za-z]\\d+)");

    private MessageRender messageRender;
    private ItineraryService itineraryService;

    public TrainStreamResponseHandler(String toolName, String agentId, ChatEnvironmentMemory chatEnvironmentMemory, MemoryManager memoryManager, MessageRender messageRender,ItineraryService itineraryService) {
        this("", "《([^》]+)》", toolName,agentId,chatEnvironmentMemory,memoryManager,messageRender,itineraryService);
    }

    public TrainStreamResponseHandler(String questionId, String itemRegex,  String toolName,String agentId,
                                      ChatEnvironmentMemory chatEnvironmentMemory,MemoryManager memoryManager,MessageRender messageRender,ItineraryService itineraryService) {
        if(!StringUtils.hasText(itemRegex)) itemRegex = "《([^》]+)》";

        this.bracketPattern = Pattern.compile(itemRegex);
        this.questionId = questionId;
        this.toolName = toolName;
        this.agentId = agentId;
        this.chatEnvironmentMemory = chatEnvironmentMemory;
        this.memoryManager = memoryManager;
        this.messageRender = messageRender;
        this.itineraryService = itineraryService;
    }


    @Override
    public Flux<String> handle(Map<String, Object> handlerContext, Flux<String> flux) {
        // Get tool response data
        String message = (String) handlerContext.get(MESSAGE);
        // Get conversation ID
        String conversationId = (String) handlerContext.get(CONVERSATION_ID);
        // Convert tool response data to map
        Map<String, Object> toolResponseMap = getToolResponseMap(message);
        if (toolResponseMap == null) {
            log.error("TrainStreamResponseHandler toolResponseMap is empty，agentId:{},message:{}", agentId,message);
            String errorMsg = I18nUtil.getMessage(I18nConstants.COMMON_SYSTEM_ERROR_MESSAGE_KEY);
            memoryManager.setMessageRight(conversationId, Collections.singletonList(new AssistantMessage(errorMsg, Map.of("agentId", agentId))));
            return Flux.just(BaseAgentNewResp.sendAnswerMsg(questionId, errorMsg));
        }
        // Get card type
        String cardType = getCardTypeByToolName(toolName).getCardType();
        List<String> matchedKeys = new ArrayList<>();
        // Create a buffer accumulation stream processor
        StringBuilder buffer = new StringBuilder();
        // Store markdown formatted recommendation information
        StringBuilder markdownTrainRecommendations = new StringBuilder();
        AtomicInteger index = new AtomicInteger(0);
        return flux.flatMap(chunk -> {
            if (!StringUtils.hasText(chunk)) {
                return Flux.empty();
            }

            // Accumulate content
            buffer.append(chunk);
            String content = buffer.toString();

            // First try to match 《》pattern
            Matcher matcher = bracketPattern.matcher(content);
            Flux<String> response = processMatches(matcher, buffer, toolResponseMap, cardType, matchedKeys, message, conversationId,markdownTrainRecommendations,index);
            if (response != null) {
                return response;
            }

            // If no 《》match found, try matching <> pattern
            matcher = angleBracketPattern.matcher(content);
            response = processMatches(matcher, buffer, toolResponseMap, cardType, matchedKeys, message, conversationId,markdownTrainRecommendations,index);
            if (response != null) {
                return response;
            }

            return Flux.empty();

        }).concatWith(Mono.fromCallable(() -> {
            String remainingContent = buffer.toString();

            // Add suggestions and summary content
            if(StringUtils.hasText(remainingContent)){
                markdownTrainRecommendations.append(remainingContent.trim());
            }
            // Store markdown formatted recommendations in context
            if (!markdownTrainRecommendations.isEmpty()) {
                memoryManager.setMessageRight(conversationId, Collections.singletonList(new AssistantMessage(markdownTrainRecommendations.toString(), Map.of("agentId", agentId))));
            }
            return StringUtils.hasText(remainingContent) ?
                    BaseAgentNewResp.sendAnswerMsg(null, remainingContent) : "";
        })).filter(StringUtils::hasText);  // Filter out empty strings
    }

    /**
     * Process matching results
     * @return Returns response string if match is successful, otherwise returns null
     */
    private Flux<String> processMatches(Matcher matcher, StringBuilder buffer, Map<String, Object> toolResponseMap,
                                  String cardType, List<String> matchedKeys, String message, String conversationId, StringBuilder markdownTrainRecommendations, AtomicInteger index) {
        if (matcher.find()) {
            // Extract train number
            String key = extractKey(matcher);

            int endIndex = matcher.start();
            // Recommendation reason
            String reason = extractReason(buffer.toString(), endIndex);

            matchedKeys.add(key);

            // Update buffer, remove processed part
            buffer.delete(0, matcher.end());

            // Get trainTicket by trainNum
            SearchAllTrainTool.TrainTicket trainTicket = (SearchAllTrainTool.TrainTicket)toolResponseMap.get(key);
            if (trainTicket == null) {
                log.error("TrainStreamResponseHandler toolResponseMap.get(key) is empty，key:{},toolResponseMap:{}", key, JsonUtil.toJson(toolResponseMap));
                return null;
            }

            // Since detail query needs departure and arrival stations, store related information in environment variables
            // Only need to store request parameters (departure station, arrival station, departure time)
            SearchTrainDetailTool.Request requestParam = new SearchTrainDetailTool.Request(trainTicket.departStationName(), trainTicket.destStationName(), extractDate(trainTicket.departureTime()), null);
            // Temporarily store data in cache
            Map<String, Object> map = new HashMap<>();
            map.put(trainTicket.trainNum(), requestParam);
            chatEnvironmentMemory.add( conversationId, map);

            // Build card information
            TrainCardData trainCardData = SearchAllTrainTool.buildCardInfo(trainTicket,conversationId);
            if(trainCardData == null){
                log.error("TrainStreamResponseHandler buildCardInfo error trainTicket:{}",JsonUtil.toJson(trainTicket));
                return null;
            }
            // Record markdown information to be stored
            String saveInMdReason = !reason.contains("：") ? reason :reason.substring(reason.indexOf("：")+1);
            buildMarkdownAssistantMsg(markdownTrainRecommendations,trainCardData,saveInMdReason,conversationId);

            Map<String, Object> trainDataMap = JsonUtil.toMap(JsonUtil.toJson(trainCardData));
            trainDataMap.put("recommendationReason",saveInMdReason);
            // Markdown message
            String markdownResponse = null;
            String viewMarkdownResponse = null;
            if (CardTypeEnum.TRAIN_TYPE.getCardType().equals(cardType)) {
                long indexNumber = NumberingUtil.incrWithExpire(conversationId);
                String acceptLanguage = I18nUtil.getLocaleLanguageTag(conversationId);
                markdownResponse = messageRender.renderTemplate("train/train-list-one.ftlh",Map.of(I18nConstants.LANGUAGE_TAG, acceptLanguage), trainDataMap);
                viewMarkdownResponse = NumberingUtil.getNumberMarkdown(markdownResponse, "#### ", indexNumber);
                // Save Command
                CommandManagerUtil.saveListCommand(conversationId, cardType + "+" + indexNumber, "stringMatch" , "defaultCallAgentAction",
                        Maps.of(DefaultCallAgentAction.PARAM_AGENT_ID, TrainAgentConfig.AGENT_ID, DefaultOutputTextAction.KEY_MESSAGE, key, StringMatch.KEY_STRING, List.of(String.valueOf(indexNumber), key)));
            }
            if (matchedKeys.size() == 1) {
                BaseItinerary trainItinerary = itineraryService.getItinerary(conversationId, null);
                if(trainItinerary != null) {
                    Resource trainResource = new TrainResource(trainTicket.departureCityName(), trainTicket.arrivalCityName(), trainTicket.departureTime(),JsonUtil.toJson(trainCardData), markdownResponse, new BigDecimal(trainCardData.getPrice()), ResourceStatus.PENDING);
                    itineraryService.saveResource(conversationId, null, trainResource);
                }
            }

            // Build message to return to frontend
            String response = createItemResponse(trainCardData, viewMarkdownResponse, cardType);
            log.info("trainStreamResponseHandler handle message:{},response:{}", message, response);
            return Flux.just(response);
        }
        return null;
    }

    /**
     * Extract key from matcher
     */
    private String extractKey(Matcher matcher) {
        String fullMatch = matcher.group(1).trim();
        Matcher trainNumberMatcher = trainNumberPattern.matcher(fullMatch);
        if (trainNumberMatcher.find()) {
            return trainNumberMatcher.group(1);
        }
        // If no specific pattern matches, return original string
        return fullMatch;
    }

    /**
     * Extract recommendation reason
     */
    private String extractReason(String content, int endIndex) {
        return content.substring(0, endIndex).trim();
    }

    /**
     * Create item response
     */
    private String createItemResponse(Object data, String reason, String cardType) {
        return BaseAgentNewResp.sendCardMsg(
                cardType,
                Collections.singletonList(data),
                reason,
                questionId
        );
    }


    private Map<String, Object> getToolResponseMap(String toolResponse) {
        SearchAllTrainTool.Response response = JsonUtil.fromJson(toolResponse, SearchAllTrainTool.Response.class);
        if(response != null && response.successCode() && !CollectionUtils.isEmpty(response.data())){
            return response.data().stream().collect(Collectors.toMap(trainTicket -> {
                //id related
//                return generateNumericTrainId(trainTicket.trainNum(),trainTicket.departureTime(),trainTicket.arrivalTime());
                return trainTicket.trainNum();
            },trainTicket->trainTicket,(k1,k2)->k2));
        }
        return null;
    }

    private CardTypeEnum getCardTypeByToolName(String toolName) {
        switch (toolName) {
            case SearchAllTrainTool.TOOLNAME:
                return CardTypeEnum.TRAIN_TYPE;
            case SearchTrainDetailTool.TOOLNAME:
            default:
                return CardTypeEnum.TRAIN_DETAIL;
        }
    }


}